#!/usr/bin/python

import os
import sys
import math
from numpy import *
from numpy.random import random_integers
import logging
import getopt
import copy
################################################
# main
################################################
loggingLevel = logging.INFO
consoleLogger = None
loggingFileName = ''
module_name = ""
PRECISION = 0.1
CMD_OPTIONS = """Command line options: 
Optional -d : Enables debug logs
Optional -m [module] : Module logic to invoke  
Optional -l [logfile] : Log file name
Optional -c : Logs to console."""
GETOPT_OPTIONS      =   "m:cdl:"

"""
Returns the adjacency matrix from the specified inputStream.
Input format must comply to the following structure:
[n: Dimension of Matrix] [n^2 integers: 1 if there is an edge, 0 otherwise]
In case of incorrect input throws Exception. 
"""
def readAdjacencyMatrix(inputStream):
    
    line = (inputStream.readline()).rstrip("\r\n")
    logging.debug("Parsing adjacency matrix file contents : %s" % line)
    params =line.split(" ")
    
    # First parameter must be the dimension.
    adjMatSize = int(params[0])
    logging.debug("Matrix size is %d" % adjMatSize)
    
    # Validate the number of elements in the array.
    #logging.debug("Remaining matrix items are of size %d." % (len(params)-1))
    #assert((len(params)-1) == size**2)
    
    # Generate matrix container
    logging.debug("Creating empty array")
    iparams = []
    
    for i in range(adjMatSize):
        logging.debug("Parsing adjacency matrix file row %d." % (i+1))
        line = (inputStream.readline()).rstrip("\r\n")
        logging.debug("Parsing data [%s] for array index %d." %(line,i) )
        params =line.split(" ")
        for j in range(adjMatSize):
            iparams.append(int(params[j]))
    
    # Transform to matrix object.
    iarray = array(iparams)
    iarray.shape=adjMatSize,adjMatSize
    
    imatrix = mat(iarray.copy())
    logging.debug("Matrix is : %s" % str(imatrix) )
    
    return imatrix

"""
Reads the group list file into a python list.
"""
def readGroupList(inputStream):
    
    line = (inputStream.readline()).rstrip("\r\n")
    logging.debug("Parsing group file contents : %s" % line)
    params =line.split(" ")
    
    # First parameter must be the dimension.
    #size = int(params[0])
    #logging.debug("Matrix size is %d" % size)
    
    # Validate the number of elements in the array.
    #logging.debug("Remaining matrix items are of size %d." % (len(params)-1))
    #assert((len(params)-1) == size)
    
    # Returns a python list of the remaining items, as integers
    return ([int(k) for k in params])
    
"""
Parses command line parameters. 
Returns the following tuple: matrix input file from which to read the matrix data, 
the group input file for group data.
"""
def parseCmdLineParameters(argv):
    
    global loggingLevel, consoleLogger, loggingFileName, module_name
    global matFile, groupFile , subgroupFile
    inputfile = sys.stdin
    
    
    if(len(argv) == 0 or len(argv) == 1):
        print CMD_OPTIONS
        sys.exit(1)
    if(len(argv) == 2):
        return (file(argv[0],"r"),file(argv[1],"r"))
        
    (optlist,iparameters) = getopt.getopt(argv[1:], GETOPT_OPTIONS, "")
    # assert that we did get the matrix and group input files.
    
    #assert(len(ifiles)==2)
        
    for opt in optlist:
        if(opt[0] == "-d"):
            loggingLeCMD_OPTIONSvel = logging.DEBUG
        elif (opt[0] == "-l"):
            loggingFileName = opt[1]
        elif (opt[0] == "-c"):
            # define a Handler which writes INFO messages or higher to the sys.stderr
            consoleLogger = logging.StreamHandler()
            # set a format which is simpler for console use
            formatter = logging.Formatter('%(name)-12s: %(levelname)-8s %(message)s')
            # tell the handler to use this format
            consoleLogger.setFormatter(formatter)
            
            if(logging.NOTSET == loggingLevel):
                loggingLevel = logging.INFO
        elif (opt[0] == "-m"):
            module_name = opt[1] 
        else:
            print "Unknown parameter"
    
    if module_name == "modularity_mat":
        matFile = file(iparameters[0],"r")
        groupFile = file(iparameters[1],"r")
    elif module_name == "spectral_div":
        matFile = file(iparameters[0],"r")
        groupFile = file(iparameters[1],"r")
        PRECISION = float(iparameters[2])
    elif module_name == "improve_div":
        matFile = file(iparameters[0],"r")
        groupFile = file(iparameters[1],"r")
        subgroupFile = file(iparameters[2],"r")
    else:
        print "Invalid module"
        sys.exit(1)

def setLoggingConfiguration():
    
    global loggingLevel, consoleLogger, loggingFileName
    
    logging.basicConfig(level=loggingLevel,
                    format='%(asctime)s %(levelname)s %(message)s',
                    filename= loggingFileName,
                    filemode='w')
    
    if(None != consoleLogger):
        consoleLogger.setLevel(logging.INFO)
        # add the handler to the root logger
        logging.getLogger('').addHandler(consoleLogger)
    

"""
Returns a matrix of dimension 1 ( vector ) that holds the rank of each node
of adjMatrix.
"""
def computeRankVector(adjMatrix):
    
    size = adjMatrix.shape[0]
    # Replace with zeros() call.
    rankV = arange(0,size,1)
    for i in range(size):
        rankV[i] = adjMatrix[i].sum()
    return rankV

"""
Computes the f(g) vector.
"""
def compute_f_g_vector(ModularityMatrix):
    
    logging.debug("Computing f_g vector")
    
    size = ModularityMatrix.shape[0]
    
    # TODO Replace with zeros() call.
    fgV = arange(0.0,size,1)
    
    for i in range(size):
        fgV[i] = float(ModularityMatrix[i].sum())
    
    logging.debug("f_g vector is %s" % str(fgV))
    return fgV

""" 
Returns a new matrix. Dimensions are the same as adjMatrix.
The contents are the modularity matrix of adjMatrix.
"""
def computeModularityMatrix(adjMatrix):
    
    M = float64(adjMatrix.sum())
    K = computeRankVector(adjMatrix)
    
    # Transform to matrix object.
    zarray = zeros(adjMatrix.shape[0]**2, dtype=float64)
    zarray.shape = adjMatrix.shape
    
    modMat = mat(zarray.copy())
    
    for i in range(modMat.shape[0]):
        for j in range(modMat.shape[0]):
            modMat[i,j] = float64(float64(adjMatrix[i,j]) - (float64(K[i]*K[j]/M)))
            logging.debug("Setting for indexes [%d][%d] the value [%lf]" % (i,j, float(modMat[i,j]))) 
    
    return modMat

"""
Returns a sub-matrix of origmat. Indexes to be used are to be passed in grp.
"""
def subclassMatrixByGroup(origmat, grp):
    
    sublist = []
    
    for i in grp:
        for j in grp:
            sublist.append(float64(origmat[i,j]))
            
    subarray = array(sublist)
    subarray.shape=len(grp),len(grp)
    
    subMatrix = mat(subarray.copy())
    
    return subMatrix


"""
If i == j, return 1, else 0.
"""
def delta(i,j):
    if(i==j):
        return 1
    
    return 0

"""
Returns the Matrix B^[g] according to the specification in the paper.
"""
def compute_B_g_tri(A, B):
    
    logging.debug("Computing B^[g]")
    V = compute_f_g_vector(B)
    sublist = []
    size = A.shape[0]
    for i in range(size):
        for j in range(size):
            sublist.append( B[i,j] - delta(i, j)*V[i])
    
    subarray = array(sublist)
    subarray.shape=size,size
    
    B_tri = mat(subarray.copy())
    
    return B_tri

"""
Computes the leading eigenpair of the matrix A.
Returns (eigenvalue,eigenvector).
"""
def computeLeadingEigenPair(A):
    
    PRECISION = 0.0001
    
    size = A.shape[0]
    X = arange(0.0,size ,1)
    
    AA = asarray(A)
    oldit = dot(AA,X)
    
    # normalize here?
    newit = dot(AA,oldit)
    
    eigenVal = float(dot(newit, oldit)) / float(dot(oldit,oldit))
    
    stop_cond = geom_len(newit - oldit*eigenVal)  < PRECISION  * geom_len(oldit)
    
    
    while(stop_cond == False):
        
        newit = newit/(abs(newit).max())
        oldit = newit;
        
        newit = dot(AA,newit)
    
        eigenVal = dot(newit, oldit) / dot(oldit,oldit)
        
        l1 = geom_len(newit - oldit*eigenVal)  
        l2 = PRECISION  * geom_len(oldit)
        stop_cond =  l1 < l2   
    
    
    # Note that newit is (A^m * X ).
    localMin = newit.min()
    newit = newit / localMin
    return (eigenVal,newit)


"""
Computes the leading eigenpair using the FAST method.
"""
def fastcomputeLeadingEigenPair(A,B,grp, PRECISION , M):
    
    subA = subclassMatrixByGroup(A, G)
    subB = subclassMatrixByGroup(B, G)
   
    size = subA.shape[0]
    
    # 1 . Get B^[g]
    subB_tri = compute_B_g_tri(subA,subB)
    
    f_gV = compute_f_g_vector(subB)
    #Mg = float(sum(subA))
    Mg = M
    
    # 2. prepare D
    D = zeros(size*size).reshape(size,size)
    
    for i in range(size):
        D[i,i] = f_gV[i]
    
    # 3 . Prepare k_g
    subKV = computeRankVector(subA)              
    
    # 4. ||A||1 == coefA1
    coefA1 =  abs(subB_tri).cumsum(axis = 0).max()
    
    X = arange(0.0,size,1)
    oldit = X
    
    oldit = sparse_calculation(X, subA, subKV, Mg, coefA1, D)
    
    newit = sparse_calculation(oldit, subA, subKV, Mg, coefA1, D)
    
    eigenVal = dot(newit, oldit) / dot(oldit,oldit)
    
    stop_cond = geom_len(newit - oldit*eigenVal)  < PRECISION  * geom_len(oldit)
    print "first " , stop_cond, newit

    while(stop_cond == False):
        
        newit = newit/(abs(newit).max())
        oldit = newit;
        
        newit = sparse_calculation(oldit, subA, subKV, Mg, coefA1, D)
        
        eigenVal = dot(newit, oldit) / dot(oldit,oldit)
        
        l1 = geom_len(newit - oldit*eigenVal)  
        l2 = PRECISION  * geom_len(oldit)
        stop_cond =  l1 < l2   
    
    
    # Note that newit is (A^m * X ).
    
    newit = newit/(abs(newit).max())
    return (eigenVal,newit)

def sparse_calculation(X,subA,subkV,M,coef,subD):
    
    AA = asarray(subA)
        
    # calculation!
    tmp1 = dot(AA,X)
    tmp2 = dot(subkV.transpose(),X)/float(M) * subkV
    tmp3 = dot(subD,X)
    tmp4 = coef*X
    
    Y = tmp1 - tmp2 - tmp3 + tmp4
    return Y
    
"""
Returns the length of the vector V calculates as sqrt(V*V) 
"""
def geom_len(V):
    return sqrt(dot(V,V))

"""
Prints matrix according to specification.
"""
def printfloatMatrix(M):
    
    size = M.shape[0] 
    
    print "%d \n" %size,
    
    for i in range(size):
        for j in range(size):
            sys.stdout.write( "%lf " % float(M[i,j]))
        if( i+1 != size):
            sys.stdout.write("\n")
        
    
"""
Implementation of the modularity_mat exercise.
"""
def modularitymat(A,G):
    
    # Core Matrix   
    B = computeModularityMatrix(A)
    
    subA = subclassMatrixByGroup(A, G)
    subB = subclassMatrixByGroup(B, G)
   
#    printfloatMatrix(subB)
    
    B_tri = compute_B_g_tri(subA,subB)
    printfloatMatrix(B_tri)

"""
Implementation of the spectral_div exercise.
"""
def spectral_div(A,G):
    
    B = computeModularityMatrix(A)
    
    M = float64(A.sum())
    subA = subclassMatrixByGroup(A, G)
    subB = subclassMatrixByGroup(B, G)
    
    subB_tri = compute_B_g_tri(subA,subB)
    coefA1 =  abs(subB_tri).cumsum(axis = 0).max() 
    shiftedM = shiftmatrix(subB_tri, coefA1)
    eigenpair1 =computeLeadingEigenPair(shiftedM)
    
    eigenpair2 = fastcomputeLeadingEigenPair(A,B,G,PRECISION,M)   
    
    s = None
    
    newsplit = []
    remaining = []
    dQ = 0.0
    (eigenvalue, eigenvector) = eigenpair2 
    
    # Subtracting the matrix shift addition. 
    eigenvalue -=coefA1
    
    if(eigenvalue <= 0):
        newsplit= G
    else:
        size = len(G)
    
        s = zeros(size, dtype=int)
        for i in range(size):
            if (eigenvector[i] > 0):
                s[i] =1
            else:
                s[i] = -1
    
        s = asarray(s)
        dQ = 0.5*dot(dot(s.transpose(), subB_tri),s)
    
        
        for i in range(len(G)):
            if(s[i] == 1):
                newsplit.append(G[i])
            else:
                remaining.append(G[i])
    
    
    # Output!
    sys.stdout.write("%lf" % dQ)
    sys.stdout.flush()
    
    if(len(newsplit) == 0):
        printIntVector(remaining)
    elif(len(remaining) ==0 ):    
        printIntVector(newsplit)
    elif(newsplit[0] < remaining[0]):
        printIntVector( newsplit)
        printIntVector (remaining)
    else:
        printIntVector( remaining)
        printIntVector( newsplit)
    
    
    newsplit = asarray(newsplit)
    remaining = asarray(remaining)
    
    return (dQ, newsplit, remaining)

"""
Utility function to print a vector of integers.
"""
def printIntVector(V):
    sys.stdout.write("\r\n")
    for i in range(len(V)):
        sys.stdout.write("%d " % V[i])
    sys.stdout.flush()
        
"""
Returns a new Matrix Z` = Z + I*coef
"""    
def shiftmatrix(Z,coef):
    size = Z.shape[0]
    Q = asarray(zeros(size*size)).reshape(size,size)
    for i in range(size):
        for j in range(size):
            Q[i,j]= Z[i,j]
            if(j==i):
                Q[i,i] += coef
        
    return Q

"""
Computes the improve_div exercise.
Parameters: Adjacency matrix A, Group G, division s.
Returns 
"""
def compute_basic_improve_div(A,G,s):
    gCounter = 0
    #print("Computing modmat");
    B = computeModularityMatrix(A)
    #print("After modmat");
    M = float64(A.sum())
    subA = subclassMatrixByGroup(A, G)
    subB = subclassMatrixByGroup(B, G)
    
    for x1 in range(subB.shape[0]):
        print "[",
        for x2 in range(subB.shape[0]):
            print "%lf " % subB[x1,x2],
        print "]"
    subB_tri = compute_B_g_tri(subA,subB)
    
    s = zeros(len(G))
    s = []
    for elem in G:
        if(elem in g1):
            s.append(float64(1.0))
        else:
            s.append(float64(-1.0))
            
    s = asarray(s)
    score = asarray([float64(-sys.maxint) for elem in range(len(G))])
    # @TODO replace with MAX_DOUBLe or something
    dQ = 1.0 
    q0 = 0.0
    origmodularity = 0.5*dot(dot(s.transpose(), subB_tri),s)
    
    indices = zeros(len(G), dtype=float64)
    improve = zeros(len(G), dtype=float64)
    while(dQ >0.0001):
    
        Unmoved = copy.copy(G)
        glist = list(G)
        # trying to find an improvement of the partition defined by s
        
        for i in range(len(G)):
            q0 = 0.5*dot(dot(s.transpose(), subB_tri),s)[0,0]
            
            # @todo: Make score use the same indexes as Unmoved
            score = asarray([float64(-sys.maxint) for elem in range(len(G))])
            for k in Unmoved:
                kIndex = glist.index(k)
                s[kIndex] = -s[kIndex]
                score[kIndex] = 0.5*dot(dot(s.transpose(), subB_tri),s)[0,0] - q0
               # if(abs(score[kIndex]) < 0.001):
               #    score[kIndex] = 0
                s[kIndex] = -s[kIndex]
               
            
            # Moving vertex j0 with a maximal score
            lscore = list(score)
            
            #print score
            print "[",
            for x in lscore:
                if(x==(-sys.maxint)):
                    print "INT_MIN ",
                else:
                    print "%lf " %x ,
            print "]" 
            lscore.reverse()
            jtag = lscore.index(max(lscore))
            lscore.reverse()
            jtag = len(score) - jtag -1
        #    print "jtag : %d " %jtag
            s[jtag] = -s[jtag]
            indices[i] = jtag
            
            if(i==0):
               improve[i] = score[jtag]
            else:
                improve[i] = improve[i-1] + score[jtag]
    
            Unmoved.remove(G[jtag])
            gCounter= gCounter+1
            
        
        # Find the maximum improvement of s and update s accordingly
        limprove = list(improve)
        limprove.reverse()
        itag = limprove.index(max(limprove))
        limprove.reverse()
        itag = len(limprove) - itag -1
        print "Max improv [%d,%lf]" %(itag,limprove[itag])
        
        r1 = range(itag+1, len(G))
        r1.reverse()
        for i in r1:
            j = indices[i]
            s[j] = -s[j]
            
        if(itag == len(G) -1 ):
            dQ = 0
        else:
            dQ = improve[itag]
            
        # End of while loop
    dQ = 0.5*dot(dot(s.transpose(), subB_tri),s)
    return (dQ-origmodularity,s)

def improve_div(A,G,g1):
    
            
    (dQ,sdiv) = compute_basic_improve_div(A, G, g1)
    
    # print output
    newsplit = []
    remaining = []
    
    for i in range(len(G)):
        if(sdiv[i] == 1):
            newsplit.append(G[i])
        else:
            remaining.append(G[i])
    
    
    # Output!
    sys.stdout.write("%lf" % dQ)
    sys.stdout.flush()
    
    if(len(newsplit) == 0):
        printIntVector(remaining)
    elif(len(remaining) ==0 ):    
        printIntVector(newsplit)
    elif(newsplit[0] < remaining[0]):
        printIntVector( newsplit)
        printIntVector (remaining)
    else:
        printIntVector( remaining)
        printIntVector( newsplit)
    
    
#####################################################################
# Main

if __name__ == '__main__':
    
    # Initialization routines
    parseCmdLineParameters(sys.argv)
    
    setLoggingConfiguration()
    
    # Reads the input from the designated file.
    A = readAdjacencyMatrix(matFile)
    logging.debug( "Adjacency matrix:\n" + str(A) )
    
    G = readGroupList(groupFile)
    
    if(module_name == "modularity_mat"):
        modularitymat(A, G)
    elif (module_name == "spectral_div"):
        spectral_div(A, G)
    elif (module_name == "improve_div"):
        g1  = readGroupList(subgroupFile)
        improve_div(A,G,g1)
       
        
         
    